﻿Imports System
Imports System.Collections.Generic
Imports System.Threading

Namespace System.Workflow.Runtime
    Friend NotInheritable Class WorkBatchCollection
        Inherits Dictionary([Of] Object, WorkBatch)
        ' Methods
        Friend Sub ClearSubBatches()
            SyncLock Me.mutex
                Dim batch As WorkBatch
                For Each batch In MyBase.Values
                    batch.Dispose
                Next
                MyBase.Clear
            End SyncLock
        End Sub

        Friend Sub ClearTransientBatch()
            Me.RollbackBatch(Me.transientBatchID)
        End Sub

        Private Function FindBatch(ByVal id As Object) As WorkBatch
            Dim batch As WorkBatch = Nothing
            SyncLock Me.mutex
                MyBase.TryGetValue(id, batch)
            End SyncLock
            Return batch
        End Function

        Friend Function GetBatch(ByVal id As Object) As IWorkBatch
            Dim batch As WorkBatch = Nothing
            SyncLock Me.mutex
                If MyBase.TryGetValue(id, batch) Then
                    Return batch
                End If
                batch = New WorkBatch(Me)
                MyBase.Add(id, batch)
            End SyncLock
            Return batch
        End Function

        Friend Function GetMergedBatch() As WorkBatch
            SyncLock Me.mutex
                Dim batch As New WorkBatch(Me)
                Dim batch2 As WorkBatch
                For Each batch2 In MyBase.Values
                    batch.Merge(batch2)
                Next
                Return batch
            End SyncLock
        End Function

        Friend Function GetNextWorkItemOrderId(ByVal pendingWork As IPendingWork) As Long
            Return Interlocked.Increment((Me._workItemOrderId))
        End Function

        Friend Function GetTransientBatch() As IWorkBatch
            Return Me.GetBatch(Me.transientBatchID)
        End Function

        Friend Sub RollbackAllBatchedWork()
            SyncLock Me.mutex
                Dim batch As WorkBatch
                For Each batch In MyBase.Values
                    batch.Complete(False)
                    batch.Dispose
                Next
                MyBase.Clear
            End SyncLock
        End Sub

        Friend Sub RollbackBatch(ByVal id As Object)
            SyncLock Me.mutex
                Dim batch As WorkBatch = Me.FindBatch(id)
                If (Not batch Is Nothing) Then
                    batch.Complete(False)
                    batch.Dispose
                    MyBase.Remove(id)
                End If
            End SyncLock
        End Sub


        ' Properties
        Friend Property WorkItemOrderId As Long
            Get
                Return Interlocked.Read((Me._workItemOrderId))
            End Get
            Set(ByVal value As Long)
                SyncLock Me.mutex
                    Me._workItemOrderId = value
                End SyncLock
            End Set
        End Property


        ' Fields
        Private _workItemOrderId As Long
        Private mutex As Object = New Object
        Private transientBatchID As Object = New Object
    End Class
End Namespace

